#include <hwt.h>
#include "encode.h"
#include <bitstream.h>
#include "embed.h"
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <huffman/huffman.h>
//#include <libjbig/jbig.h>
#include <jbig.h>
#include <md5/md5.h>

#define PAYLOAD_MESSAGE_LENGTH 16065 - 4 //bytes

struct bie_t {
  size_t size;
  unsigned char *buffer;
};

void output_bie(unsigned char *start, size_t len, void *out) {
  struct bie_t *b = (struct bie_t *)out;
  memcpy(b->buffer + b->size, start, len);
  b->size += len;
}

int tian_grayscale_encode(
			  unsigned char **marked,
			  unsigned char **data,
			  unsigned int height,
			  unsigned int width,
			  unsigned int threshold,
			  const char *payload_fname,
			  const int verbose
			  )
{

#ifdef DEBUG
  int checksum, checksum_, j;
#endif

  int i;
  unsigned char *lm, **average, *lsb;
  int **differences;
  difference_class **membership;
  class_count cc;

#ifdef DEBUG
  checksum = 0;
  for(i=0; i<height; ++i)
    for(j=0; j<width; ++j)
      checksum += data[i][j];
  printf("+ original bitmap checksum: %u\n", checksum);
#endif

  /*
    0. compute the payload
  */
  const unsigned int payload_message_length = payload_fname ? PAYLOAD_MESSAGE_LENGTH : 16 /*128-bit md5-digest*/; //bytes
  unsigned char payload[4 + payload_message_length];

  if(payload_fname == NULL) {
    //md5 checksum
    if(verbose) printf("> computing the payload (128-bit md5 digest)\n");
    md5_state_t md5_state;
    md5_init(&md5_state);
    for(i=0; i<height; ++i)
      md5_append(&md5_state, data[i], width);
    md5_finish(&md5_state, 4 + payload);
  }
  else {
    //read from file
    FILE *f = fopen(payload_fname, "rt");
    fgets((char *)payload + 4, PAYLOAD_MESSAGE_LENGTH, f);
  }

  //length header (32 bit as 8-bit words)
  payload[3] = payload_message_length & 255; //8 LSB
  payload[2] = (payload_message_length >> 8) & 255;
  payload[1] = (payload_message_length >> 16) & 255;
  payload[0] = (payload_message_length >> 24) & 255;

#ifdef DEBUG
  unsigned char *digest = 4 + payload;
  checksum = 0;
  for(i=0; i<payload_message_length; ++i)
    checksum += digest[i];
  printf("+ payload-message checksum: %u\n", checksum);
#endif
  
  //convert to FLAT bitstream: 1 bit per 8-bit word
  unsigned int length_P;
  unsigned char *pbs = to_bitstream_flat(payload, 4 + payload_message_length, &length_P);

#ifdef DEBUG
  printf("|P| = %d bits (32-bit header + 128-bit md5 digest)\n", length_P);
#endif
  if(verbose) printf("payload-message length: %u bytes (%u bits)\n", payload_message_length, payload_message_length * 8);


  /*
    1. compute difference-bitmap and average-bitmap (integer Haar wavelet transform)
  */
  if(verbose) printf("> computing integer Haar wavelet transform\n");
  
  //allocate the difference-bitmap
  differences = (int **)malloc(height * sizeof(int *));
  for(i=0; i<height; ++i)
    differences[i] = (int *)malloc(width/2 * sizeof(int));

  //allocate the average-bitmap
  average = (unsigned char **)malloc(height * sizeof(unsigned char *));
  for(i=0; i<height; ++i)
    average[i] = (unsigned char *)malloc(width/2 * sizeof(unsigned char));

  hwt_bitmap(differences, average, data, height, width);



#ifdef DEBUG
  checksum_ = 0;
  for(i=0; i<height; ++i)
    for(j=0; j<width/2; ++j)
      checksum_ += differences[i][j];
  printf("+ differences checksum: %d\n", checksum_);
#endif


  /*
    2. compute membership-bitmap
  */
  if(verbose) printf("> classification (threshold: %u)\n", threshold);
  
  //allocate the membership-bitmap
  membership = (difference_class **)malloc(height * sizeof(difference_class *));
  for(i=0; i<height; ++i)
    membership[i] = (difference_class *)malloc(width/2 * sizeof(difference_class));

  diff_classification_all(membership, &cc, differences, average, height, width/2, threshold);

  unsigned int total_embedding_capacity = cc.n_EZ + cc.n_EN1 + cc.n_EN2 + cc.n_CH;

#ifdef DEBUG
  //print embedding capacity etc.
  printf("|EZ| = %d bits\n", cc.n_EZ);
  printf("|EN1| = %d bits\n", cc.n_EN1);
  printf("|EN2| = %d bits\n", cc.n_EN2);
  printf("|CH| = %d bits (|CHX| = %d bits)\n", cc.n_CH, cc.n_CHX);
  printf("|NC| = %d bits\n", cc.n_NC);
  printf("total embedding capacity: %d bits (on %d differences)\n", total_embedding_capacity, width * height / 2);
  printf("expansion capacity: %u bits\n", cc.n_EZ + cc.n_EN1 + cc.n_CHX);
#endif


  /*
    3. compute (and compress) location-map (LM)
    format:
    | HDR | w1 ... wN |
    with:
    - HDR (int, 32 bit): n. bits in compressed LM (including HDR)
    - w1 ... wN (8-bit words): compressed LM
  */
  if(verbose) printf("> building location-map\n");

  //allocate the location-map
  unsigned int lm_length = height * width/2;
  lm = (unsigned char *)malloc(lm_length * sizeof(unsigned char));
  diff_map(lm, membership, height, width/2);

  //convert to COMPACT bitstream (8-bit words, 0-padded)
  unsigned int lmbs_size;
  unsigned char *lmbs = to_bitstream_compact(lm, lm_length, &lmbs_size);

#ifdef DEBUG
  checksum = 0;
  for(i=0; i<lm_length; ++i)
    checksum += lm[i];
  printf("+ location-map checksum: %d\n", checksum);
#endif

  // 3'. compress location-map
  //if(verbose) printf("> compress (Huffmann) location-map\n");
  //unsigned char *clmbs_ = (unsigned char *)malloc((4 + lmbs_words) * sizeof(unsigned char));
  //unsigned int clmbs_words = Huffman_Compress(lmbs, 4 + clmbs_, lmbs_words);
  /* unsigned int length_L = 4 + clmbs_words; */
  /* //length header (32 bit as 8-bit words) */
  /* clmbs_[3] = length_L & 255; //8 LSB */
  /* clmbs_[2] = (length_L >> 8) & 255; */
  /* clmbs_[1] = (length_L >> 16) & 255; */
  /* clmbs_[0] = (length_L >> 24) & 255; */
  //convert to FLAT bitstream
  //unsigned char *clmbs = to_bitstream_flat(clmbs_, 4 + clmbs_words, &length_L);


  if(verbose) printf("> compress (JBIG1) location-map\n");
  unsigned char *clmbs_ = (unsigned char *)malloc((4 + lmbs_size) * sizeof(unsigned char)); //enc. buffer
  struct bie_t b;
  b.size = 0;
  b.buffer = clmbs_ + 4;
  struct jbg_enc_state se;
  jbg_enc_init(&se, width/2, height, 1, &lmbs, output_bie, &b);
  jbg_enc_out(&se);
  jbg_enc_free(&se);
  size_t clmbs_size = b.size;
  unsigned int L_size = 4 + clmbs_size;
  //length header (32-bit integer)
  clmbs_[3] = L_size & 255;
  clmbs_[2] = (L_size >> 8) & 255;
  clmbs_[1] = (L_size >> 16) & 255;
  clmbs_[0] = (L_size >> 24) & 255;
  //convert to FLAT bitstream
  unsigned int length_L;
  unsigned char *clmbs = to_bitstream_flat(clmbs_, L_size, &length_L);

#ifdef DEBUG
  printf("(uncompressed) location-map size: %u bits\n", lm_length);
  //printf("Huffmann-compressed location-map size: %d bits (%f:1 ratio)\n", clmbs_words * 8, (float)lm_size / (clmbs_words*8));
  printf("JBIG-compressed location-map size: %zu bits (%f:1 ratio)\n", clmbs_size * 8, (float)lm_length / (clmbs_size * 8));
  //printf("[Huffmann] |L|: %d bits (32-bit header + %u-bit data)\n", length_L, clmbs_words);
  printf("|L|: %d bits (32-bit header + %zu-bit data)\n", length_L, clmbs_size * 8);
#endif

  //destroy location-map, uncompressed lm bitstream and compact compressed lm bitstream
  free(lm);
  free(lmbs);
  free(clmbs_);


  /*
    4. collect LSB of differences in (EN2 U CH) \ CHX
  */
  if(verbose) printf("> collecting LSB of differences in EN2 and CH\n");
  unsigned int length_C = cc.n_EN2 + cc.n_CH - cc.n_CHX;
#ifdef DEBUG
  printf("|C| = %d bits\n", length_C);
#endif
  lsb = (unsigned char *)malloc(length_C * sizeof(unsigned char));
  collect_LSB(lsb, differences, membership, height, width/2);


  
  if(verbose) printf("> check embedding capacity\n");
  unsigned int lp = length_L + length_P;
  unsigned int ub = cc.n_EZ + cc.n_EN1 + cc.n_CHX;
#ifdef CAPACITY
  fprintf(stdout, "T\t|EN1|\t|EN2|\t|L|\tratio\tmax payload\tmax bitrate\n");
  fprintf(stdout, "%u\t%u\t%u\t%u\t%.2f\t%ld bits\t%.4f\n",
	  threshold,
	  cc.n_EN1,
	  cc.n_EN2,
	  length_L,
	  (float)height * width / 2 / length_L,
	  (long)ub - length_L,
	  (float)((long)ub - length_L) / (height * width)
	  );
#endif
#ifdef DEBUG
  printf("condition: |L| + |P| <= |EZ| + |EN1| + |CHX|\n");
  printf("check: |%u| + |%u| = %u <= |%u| + |%u| + |%u| = %u\n",
	 length_L, length_P, lp, cc.n_EZ, cc.n_EN1, cc.n_CHX, ub);
#endif
  if(lp > ub) {
    printf("insufficient embedding capacity: try to increase the threshold\n");
    printf("> stop\n");
    //clean-up
    for(i=0; i<height; ++i) {
      free(differences[i]);
      free(membership[i]);
      free(average[i]);
    }
    free(differences);
    free(membership);
    free(average);
    free(clmbs);
    free(lsb);
    exit(1);
  }
#ifdef DEBUG
  else
    printf("ok\n");
#endif

#ifdef DEBUG
  printf("bitstream length: |L| + |C| + |P| = %d bits\n", length_L + length_C + length_P);
  printf("bitrate: %f\n", (float)payload_message_length * 8 / (height * width));
#endif


  /*
    5. embed
  */
  //embed by difference expansion/changing
  if(verbose) printf("> embedding\n");

  unsigned int start_row = 0, start_column = 0, embedded = 0;
  embed(differences, &start_row, &start_column, clmbs, length_L, membership, width/2, &embedded); //L
#ifdef DEBUG
  printf("embedded %u / %u bits\n", embedded, total_embedding_capacity);
#endif
  embed(differences, &start_row, &start_column, lsb, length_C, membership, width/2, &embedded); //C
#ifdef DEBUG
  printf("embedded %u / %u bits\n", embedded, total_embedding_capacity);
#endif
  embed(differences, &start_row, &start_column, pbs, length_P, membership, width/2, &embedded); //P
#ifdef DEBUG
  printf("embedded %u / %u bits\n", embedded, total_embedding_capacity);
#endif
  int pad_size = total_embedding_capacity - (length_L + length_C + length_P);
  unsigned char *pad_bs = (unsigned char *)malloc(pad_size * sizeof(unsigned char));
  for(i=0; i<pad_size; ++i)
    pad_bs[i] = 0;
#ifdef DEBUG
  printf("0-pad length: %u\n", pad_size);
#endif
  embed(differences, &start_row, &start_column, pad_bs, pad_size, membership, width/2, &embedded); //0-pad

#ifdef DEBUG
  printf("embedded bits: %u\n", embedded);
#endif

#ifdef DEBUG
  checksum_ = 0;
  checksum = 0;
  for(i=0; i<height; ++i)
    for(j=0; j<width/2; ++j) {
      checksum += average[i][j];
      checksum_ += differences[i][j];
    }
  printf("+ emb. differences checksum: %d\n", checksum_);
  printf("+ emb. averages checksum: %d\n", checksum);
#endif

  //inverse Haar transform
  if(verbose) printf("> inverse transform\n");
  ihwt_bitmap(marked, differences, average, height, width);

  /* //check */
/*   /\* anti = (unsigned char **)malloc(height * sizeof(unsigned char *)); *\/ */
/*   /\* for(i=0; i<height; ++i) *\/ */
/*   /\*   anti[i] = (unsigned char *)malloc(width * sizeof(int)); *\/ */
/*   /\* ihwt_bitmap(anti, differences, average, height, width); *\/ */
/*   //diffs */
/*   int **differences_ = (int **)malloc(height * sizeof(int *)); */
/*   for(i=0; i<height; ++i) */
/*     differences_[i] = (int *)malloc(width/2 * sizeof(int)); */
/*   //allocate the average-bitmap */
/*   unsigned char **average_ = (unsigned char **)malloc(height * sizeof(unsigned char *)); */
/*   for(i=0; i<height; ++i) */
/*     average_[i] = (unsigned char *)malloc(width/2 * sizeof(unsigned char)); */
/*   hwt_bitmap(differences_, average_, marked, height, width); */
/* #ifdef DEBUG */
/*   checksum_ = 0; */
/*   for(i=0; i<height; ++i) */
/*     for(j=0; j<width/2; ++j) */
/*       checksum_ += differences_[i][j]; */
/*   printf("+ haar-check checksum: %d\n", checksum_); */
/*   for(i=0; i<height; ++i) */
/*     for(j=0; j<width/2; ++j) */
/*       if(differences[i][j] != differences_[i][j]) { */
/* 	printf("(%d,%d): %d vs %d\n", i, j, differences[i][j], differences_[i][j]); */
/* 	printf("x = %u, y = %u\n", marked[i][2*j], marked[i][2*j+1]); */
/* 	printf("h = %d, l = %u\n", differences[i][j], average[i][j]); */
/*       } */
/* #endif */

  //clean-up
  for(i=0; i<height; ++i) {
    free(differences[i]);
    free(membership[i]);
    free(average[i]);
  }
  free(differences);
  free(membership);
  free(average);
  free(clmbs);
  free(lsb);
  free(pad_bs);
  //free payload

  return 0;
}
